# -*- coding: utf-8 -*-
"""
基于Web界面的银行客户查询助手 - 增强版
"""
import os
import asyncio
from typing import Optional
import dashscope
from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
import pandas as pd
from sqlalchemy import create_engine
from qwen_agent.tools.base import BaseTool, register_tool
import matplotlib.pyplot as plt
import io
import base64
import time
import numpy as np

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS']  # 优先使用的中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义资源文件根目录
ROOT_RESOURCE = os.path.join(os.path.dirname(__file__), 'resource')

# 配置数据库连接
MYSQL_USER = os.getenv('MYSQL_USER', 'root')
MYSQL_PASSWORD = os.getenv('MYSQL_PASSWORD', '')
MYSQL_HOST = 'localhost'
MYSQL_PORT = '3306'
DATABASE_NAME = 'bank'  # 数据库名更改为 bank

def create_connection():
    """创建数据库连接"""
    try:
        engine = create_engine(
            f'mysql+mysqlconnector://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{DATABASE_NAME}?charset=utf8mb4',
            connect_args={'connect_timeout': 10}, pool_size=10, max_overflow=20
        )
        return engine
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        return None

# ====== 更新助手的 system_prompt，添加开场白 ======
system_prompt = """
您好！欢迎使用百万客群经营助手。

我是您的智能助手，专注于帮助您高效管理和分析高净值客户数据。以下是我的核心功能：
- 查询客户基本信息（姓名、年龄、职业等）
- 分析客户的资产分布（存款、理财、基金等）
- 查看客户的交易行为和APP使用习惯
- 生成可视化图表以辅助决策
- 使用逻辑回归模型预测客户未来N个月资产提升至100万+的概率

以下是数据库中的表结构说明：

客户基础信息表结构如下：
CREATE TABLE customer_base (
    customer_id VARCHAR(32) PRIMARY KEY COMMENT '客户ID',
    name VARCHAR(100) COMMENT '客户姓名',
    age INT COMMENT '年龄',
    gender VARCHAR(10) COMMENT '性别',
    occupation VARCHAR(100) COMMENT '职业',
    occupation_type VARCHAR(50) COMMENT '职业类型标签（如：企业高管/互联网从业者/私营业主）',
    monthly_income DECIMAL(12,2) COMMENT '月收入',
    open_account_date VARCHAR(10) COMMENT '开户日期',
    lifecycle_stage VARCHAR(50) COMMENT '客户生命周期',
    marriage_status VARCHAR(20) COMMENT '婚姻状态',
    city_level VARCHAR(20) COMMENT '城市等级（一线/二线城市）',
    branch_name VARCHAR(100) COMMENT '开户网点'
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户基础信息表';

客户行为资产表结构如下：
CREATE TABLE customer_behavior_assets (
    id VARCHAR(32) PRIMARY KEY COMMENT '主键ID',
    customer_id VARCHAR(32) COMMENT '客户ID（关联customer_base表）',
    total_assets DECIMAL(16,2) COMMENT '总资产',
    deposit_balance DECIMAL(16,2) COMMENT '存款余额',
    financial_balance DECIMAL(16,2) COMMENT '理财余额',
    fund_balance DECIMAL(16,2) COMMENT '基金余额',
    insurance_balance DECIMAL(16,2) COMMENT '保险余额',
    asset_level VARCHAR(20) COMMENT '资产分层（50万以下、50-80万、80-100万、100万+）',
    deposit_flag TINYINT COMMENT '是否持有存款（1是0否）',
    financial_flag TINYINT COMMENT '是否持有理财（1是0否）',
    fund_flag TINYINT COMMENT '是否持有基金（1是0否）',
    insurance_flag TINYINT COMMENT '是否持有保险（1是0否）',
    product_count INT COMMENT '持有产品数量',
    financial_repurchase_count INT COMMENT '近1年理财复购次数',
    credit_card_monthly_expense DECIMAL(12,2) COMMENT '信用卡月均消费',
    investment_monthly_count INT COMMENT '月均投资交易次数',
    app_login_count INT COMMENT 'APP月均登录次数',
    app_financial_view_time INT COMMENT '理财页面月均停留时长(秒)',
    app_product_compare_count INT COMMENT '产品对比点击次数',
    last_app_login_time VARCHAR(19) NULL COMMENT '最近APP登录时间',
    last_contact_time VARCHAR(19) NULL COMMENT '最近联系时间',
    contact_result VARCHAR(50) COMMENT '联系结果',
    marketing_cool_period VARCHAR(10) COMMENT '营销冷却期（30天）',
    stat_month VARCHAR(7) COMMENT '统计月份（YYYY-MM）',
    FOREIGN KEY (customer_id) REFERENCES customer_base(customer_id),
    UNIQUE KEY uk_customer_month (customer_id, stat_month)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户行为资产表';

每当 exc_sql 工具返回 markdown 表格和图片时，你必须原样输出工具返回的全部内容（包括图片），不要只总结表格，也不要省略图片。这样用户才能直接看到表格和图片。

重要提醒：
1. 当工具返回包含图表的结果时，必须在你的回答中完整显示这些图表
2. 在给出分析结论后，如果有相关图表，请再次展示图表以便用户查看
3. 图表是数据分析的重要组成部分，不能省略或仅用文字描述替代
4. 确保图表链接格式正确：![图表标题](图片路径)

知识库中是以前做过的示例，如果有类似的需要进行逻辑上的参考。
"""

functions_desc = [
    {
        "name": "exc_sql",
        "description": "对于生成的SQL，进行SQL查询",
        "parameters": {
            "type": "object",
            "properties": {
                "sql_input": {
                    "type": "string",
                    "description": "生成的SQL语句",
                }
            },
            "required": ["sql_input"],
        },
    },
]

# ====== 会话隔离 DataFrame 存储 ======
# 用于存储每个会话的 DataFrame，避免多用户数据串扰
_last_df_dict = {}

def get_session_id(kwargs):
    """根据 kwargs 获取当前会话的唯一 session_id，这里用 messages 的 id"""
    messages = kwargs.get('messages')
    if messages is not None:
        return id(messages)
    return None

# ====== exc_sql 工具类实现 ======
@register_tool('exc_sql')
class ExcSQLTool(BaseTool):
    """
    SQL查询工具，执行传入的SQL语句并返回结果，并自动进行可视化。
    """
    description = '对于生成的SQL，进行SQL查询，并自动可视化'
    parameters = [{
        'name': 'sql_input',
        'type': 'string',
        'description': '生成的SQL语句',
        'required': True
    }, {
        'name': 'need_visualization',
        'type': 'boolean',
        'description': '是否需要可视化图表，默认为True',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        sql_input = args['sql_input']
        need_visualization = args.get('need_visualization', True)  # 默认需要可视化
        database = args.get('database', DATABASE_NAME)  # 使用配置的数据库名称
        engine = create_connection()
        
        try:
            df = pd.read_sql(sql_input, engine)
            
            # 检查是否是股票对比查询，如果是则不需要可视化
            if "对比" in sql_input or "比较" in sql_input:
                need_visualization = False
            
            # 检查是否是产品分布查询，强制启用可视化
            if any(keyword in sql_input.lower() for keyword in ['产品', '分布', 'flag', 'balance', '存款', '理财', '基金', '保险', '年龄', 'age']):
                need_visualization = True
            
            # 生成Markdown表格，对于年龄分布查询显示更多信息
            if 'age' in df.columns and len(df) > 10:
                # 获取全体客户数据用于年龄分布统计
                all_customers_query = "SELECT age FROM customer_base WHERE age >= 18"
                all_customers_df = pd.read_sql(all_customers_query, engine)
                
                # 年龄分布查询显示统计摘要
                age_stats = all_customers_df['age'].describe()
                md = f"年龄分布统计摘要（共{len(all_customers_df)}位客户）：\n"
                md += f"- 平均年龄：{age_stats['mean']:.1f}岁\n"
                md += f"- 年龄范围：{int(age_stats['min'])}岁 - {int(age_stats['max'])}岁\n"
                md += f"- 中位数年龄：{age_stats['50%']:.1f}岁\n"
                md += f"- 18岁以上客户：{len(all_customers_df)}人\n\n"
                
                # 显示年龄段分布
                age_counts = all_customers_df['age'].value_counts().sort_index()
                
                # 按5岁分组统计
                min_age = max(18, int(age_counts.index.min()))
                max_age = int(age_counts.index.max())
                
                age_groups = {}
                for age in range(min_age, max_age + 1):
                    group_start = (age - min_age) // 5 * 5 + min_age
                    group_end = group_start + 4
                    group_key = f"{group_start}-{group_end}岁"
                    
                    if group_key not in age_groups:
                        age_groups[group_key] = 0
                    
                    if age in age_counts.index:
                        age_groups[group_key] += age_counts[age]
                
                md += "年龄段分布：\n"
                for group, count in sorted(age_groups.items(), key=lambda x: int(x[0].split('-')[0])):
                    if count > 0:
                        md += f"- {group}：{count}人\n"
                        
            elif len(df) <= 10:
                md = df.to_markdown(index=False)
            else:
                # 显示前5行和后5行
                head_df = df.head(5)
                tail_df = df.tail(5)
                md = head_df.to_markdown(index=False) + "\n...\n" + tail_df.to_markdown(index=False)
            
            # 如果不需要可视化，直接返回表格数据
            if not need_visualization:
                # 如果是股票对比查询，添加涨跌幅计算公式和结果
                if "对比" in sql_input or "比较" in sql_input:
                    pct_chg_formula = "\n\n涨跌幅计算公式：\n涨跌幅(%) = (当日收盘价 - 昨日收盘价) / 昨日收盘价 × 100%\n\n"
                    if 'pct_chg' in df.columns:
                        pct_chg_formula += "查询结果中的涨跌幅(%)字段即为按上述公式计算的结果。\n"
                    return f"{md}{pct_chg_formula}"
                return md
            
            # 添加数据描述信息
            describe_md = "数据统计信息：\n" + df.describe().to_markdown()
            
            # 如果只有一行或没有数据，则不生成图表
            if len(df) <= 1:
                if len(df) > 0:
                    return f"{md}\n\n{describe_md}"
                return md
            
            # 自动创建目录
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            
            # 根据查询内容生成不同的图表名称
            if any(col in df.columns for col in ['deposit_flag', 'financial_flag', 'fund_flag', 'insurance_flag']):
                filename = f'product_distribution_{int(time.time()*1000)}.png'
                chart_title = '产品分布图'
            elif any(col in df.columns for col in ['deposit_balance', 'financial_balance', 'fund_balance', 'insurance_balance']):
                filename = f'balance_distribution_{int(time.time()*1000)}.png'
                chart_title = '资产分布图'
            elif 'age' in df.columns:
                filename = f'age_distribution_{int(time.time()*1000)}.png'
                chart_title = '年龄分布图'
            else:
                filename = f'data_chart_{int(time.time()*1000)}.png'
                chart_title = '数据图表'
            
            save_path = os.path.join(save_dir, filename)
            
            # 生成图表
            generate_chart_png(df, save_path)
            img_path = os.path.join('image_show', filename)
            img_md = f'![{chart_title}]({img_path})'
            
            # 确保图表路径正确并返回完整内容
            result = f"{md}\n\n{describe_md}\n\n{img_md}"
            print(f"生成图表: {img_path}")  # 调试信息
            return result
        except Exception as e:
            return f"查询执行失败: {str(e)}"

def generate_chart_png(df, save_path):
    """生成柱状图并保存"""
    # 关闭之前的图形，避免内存泄漏
    plt.close('all')
    plt.figure(figsize=(12, 8))
    
    # 检查是否包含产品分布相关的列
    product_columns = ['deposit_flag', 'financial_flag', 'fund_flag', 'insurance_flag']
    balance_columns = ['deposit_balance', 'financial_balance', 'fund_balance', 'insurance_balance']
    
    # 如果包含产品标志列，生成产品分布图
    if any(col in df.columns for col in product_columns):
        # 计算各产品的持有人数
        product_counts = {}
        product_labels = {'deposit_flag': '存款', 'financial_flag': '理财', 
                         'fund_flag': '基金', 'insurance_flag': '保险'}
        
        for col in product_columns:
            if col in df.columns:
                product_counts[product_labels[col]] = int(df[col].sum())
        
        if product_counts:
            # 创建子图布局
            has_balance = any(col in df.columns for col in balance_columns)
            if has_balance:
                plt.subplot(2, 1, 1)
            
            colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4']
            bars = plt.bar(product_counts.keys(), product_counts.values(), color=colors[:len(product_counts)])
            plt.title('客户产品持有分布', fontsize=16, fontweight='bold', pad=20)
            plt.ylabel('持有人数', fontsize=12)
            
            # 设置y轴范围，确保显示整数
            max_count = max(product_counts.values())
            plt.ylim(0, max_count * 1.2)
            
            # 设置y轴刻度为整数
            plt.gca().yaxis.set_major_locator(plt.MaxNLocator(integer=True))
            
            # 在柱子上显示数值
            for bar in bars:
                height = bar.get_height()
                plt.text(bar.get_x() + bar.get_width()/2., height + max_count*0.02,
                        f'{int(height)}人', ha='center', va='bottom', fontsize=11, fontweight='bold')
    
    # 如果包含余额列，生成余额分布图
    if any(col in df.columns for col in balance_columns):
        balance_data = {}
        balance_labels = {'deposit_balance': '存款余额', 'financial_balance': '理财余额',
                         'fund_balance': '基金余额', 'insurance_balance': '保险余额'}
        
        for col in balance_columns:
            if col in df.columns:
                balance_data[balance_labels[col]] = df[col].sum()
        
        if balance_data:
            # 如果同时有产品分布，使用子图
            has_product = any(col in df.columns for col in product_columns)
            if has_product:
                plt.subplot(2, 1, 2)
            
            colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4']
            bars = plt.bar(balance_data.keys(), balance_data.values(), color=colors[:len(balance_data)])
            plt.title('客户资产余额分布', fontsize=16, fontweight='bold', pad=20)
            plt.ylabel('余额总额（万元）', fontsize=12)
            
            # 设置y轴范围
            max_balance = max(balance_data.values())
            plt.ylim(0, max_balance * 1.2)
            
            # 在柱子上显示数值（以万元为单位）
            for bar in bars:
                height = bar.get_height()
                if height >= 10000:
                    label = f'{height/10000:.1f}万'
                else:
                    label = f'{height:.0f}'
                plt.text(bar.get_x() + bar.get_width()/2., height + max_balance*0.02,
                        label, ha='center', va='bottom', fontsize=11, fontweight='bold')
    
    # 如果既没有产品标志也没有余额数据，使用通用图表
    elif not any(col in df.columns for col in product_columns + balance_columns):
        # 检查是否是年龄分布数据
        if 'age' in df.columns:
            # 年龄分布特殊处理
            plt.clf()  # 清除当前图形内容
            plt.figure(figsize=(14, 8))  # 重新设置图形大小
            
            try:
                # 过滤18岁以上的客户
                adult_ages = df[df['age'] >= 18]['age']
                
                if len(adult_ages) == 0:
                    # 如果没有18岁以上的客户，使用全部数据
                    adult_ages = df['age']
                
                # 按年龄分组统计
                age_counts = adult_ages.value_counts().sort_index()
                
                # 创建年龄段分组（5岁一组）
                min_age = int(age_counts.index.min())
                max_age = int(age_counts.index.max())
                
                # 从18岁开始，或者从最小年龄开始
                start_age = max(18, min_age)
                
                # 创建年龄段
                age_groups = {}
                for age in range(start_age, max_age + 1):
                    group_start = (age - start_age) // 5 * 5 + start_age
                    group_end = group_start + 4
                    group_key = f"{group_start}-{group_end}岁"
                    
                    if group_key not in age_groups:
                        age_groups[group_key] = 0
                    
                    if age in age_counts.index:
                        age_groups[group_key] += age_counts[age]
                
                # 排序年龄段
                sorted_groups = sorted(age_groups.items(), key=lambda x: int(x[0].split('-')[0]))
                group_labels = [item[0] for item in sorted_groups]
                group_counts = [item[1] for item in sorted_groups]
                
                # 生成图表
                bars = plt.bar(range(len(group_labels)), group_counts, color='#4A90E2')
                plt.xticks(range(len(group_labels)), group_labels, rotation=45, ha='right')
                plt.title('客户年龄分布（18岁以上）', fontsize=16, fontweight='bold', pad=20)
                plt.ylabel('客户人数', fontsize=12)
                plt.xlabel('年龄段', fontsize=12)
                
                # 设置y轴格式，根据数据规模调整
                max_count = max(group_counts) if group_counts else 1
                
                # 根据最大值调整y轴标签格式
                if max_count >= 1000:
                    # 以千为单位显示
                    plt.gca().yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x/1000:.1f}K'))
                    plt.ylabel('客户人数（千人）', fontsize=12)
                elif max_count >= 100:
                    # 以百为单位显示
                    plt.gca().yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x/100:.1f}百'))
                    plt.ylabel('客户人数（百人）', fontsize=12)
                
                # 在柱子上显示数值
                for i, bar in enumerate(bars):
                    height = bar.get_height()
                    if height > 0:
                        if height >= 1000:
                            label = f'{height/1000:.1f}K'
                        elif height >= 100:
                            label = f'{height/100:.1f}百'
                        else:
                            label = f'{int(height)}'
                        
                        plt.text(bar.get_x() + bar.get_width()/2., height + max_count*0.01,
                                label, ha='center', va='bottom', fontsize=10, fontweight='bold')
                
                # 设置网格线，便于阅读
                plt.grid(axis='y', alpha=0.3, linestyle='--')
                        
            except Exception as e:
                print(f"年龄分布处理出错: {e}")
                # 简单回退处理
                age_counts = df['age'].value_counts().sort_index()
                ages = list(age_counts.index)
                counts = list(age_counts.values)
                
                plt.bar(ages, counts, color='skyblue')
                plt.title('客户年龄分布', fontsize=16, fontweight='bold')
                plt.ylabel('客户人数', fontsize=12)
                plt.xlabel('年龄', fontsize=12)
        
        else:
            # 其他类型的数据
            numeric_cols = df.select_dtypes(include=[np.number]).columns
            
            if len(numeric_cols) >= 2 and len(df) <= 20:  # 多列且数据不太多时使用分组柱状图
                ax = df[numeric_cols[:4]].plot(kind='bar', figsize=(12, 6), color=['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4'])
                plt.title('客户数据统计', fontsize=16, fontweight='bold')
                plt.ylabel('数值', fontsize=12)
                plt.xticks(rotation=45)
                plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
            else:
                # 单列数据或数据太多时使用简单柱状图
                if len(numeric_cols) > 0:
                    # 如果数据太多，只显示部分x轴标签
                    ax = df[numeric_cols[0]].plot(kind='bar', color='skyblue', figsize=(12, 6))
                    plt.title('数据分布', fontsize=16, fontweight='bold')
                    plt.xlabel('索引', fontsize=12)
                    plt.ylabel('值', fontsize=12)
                    
                    # 优化x轴标签显示
                    if len(df) > 20:
                        step = max(1, len(df) // 10)
                        plt.xticks(range(0, len(df), step), df.index[::step], rotation=45)
                    else:
                        plt.xticks(rotation=45)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"图表已保存到: {save_path}")  # 调试信息

# ====== 新增：逻辑回归预测工具 ======
@register_tool('lr_predict')
class LRPredictTool(BaseTool):
    """
    逻辑回归预测工具，预测客户未来N个月资产提升至100万+的概率。
    """
    description = '使用逻辑回归模型预测客户未来N个月资产提升至100万+的概率'
    parameters = [{
        'name': 'customer_id',
        'type': 'string',
        'description': '客户ID',
        'required': True
    }, {
        'name': 'prediction_period',
        'type': 'integer',
        'description': '预测周期（月数），默认为3',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        customer_id = args['customer_id']
        prediction_period = args.get('prediction_period', 3)  # 默认预测3个月
        
        # 从数据库获取客户数据
        engine = create_connection()
        if engine is None:
            return "数据库连接失败"
        
        # 获取客户历史数据
        query = f"""
        SELECT cb.name, cb.age, cb.gender, cb.monthly_income, cb.occupation, 
               cba.total_assets, cba.deposit_balance, cba.financial_balance, 
               cba.fund_balance, cba.insurance_balance, cba.product_count,
               cba.financial_repurchase_count, cba.credit_card_monthly_expense,
               cba.investment_monthly_count, cba.app_login_count,
               cba.app_financial_view_time, cba.app_product_compare_count,
               cba.stat_month
        FROM customer_base cb
        JOIN customer_behavior_assets cba ON cb.customer_id = cba.customer_id
        WHERE cb.customer_id = '{customer_id}'
        ORDER BY cba.stat_month
        """
        
        try:
            df = pd.read_sql(query, engine)
            if df.empty:
                return f"未找到客户ID为 {customer_id} 的数据"
            
            # 计算增长率
            df['prev_assets'] = df['total_assets'].shift(1)
            df['asset_growth'] = df['total_assets'] - df['prev_assets']
            df['growth_rate'] = df['asset_growth'] / df['prev_assets']
            
            # 处理无穷大和NaN值
            df['growth_rate'] = df['growth_rate'].replace([np.inf, -np.inf], np.nan)
            df['growth_rate'] = df['growth_rate'].fillna(0)
            
            # 计算平均增长率并添加到DataFrame
            avg_growth_rate = df['growth_rate'].mean()
            if np.isnan(avg_growth_rate):
                avg_growth_rate = 0
            df['avg_growth_rate'] = avg_growth_rate  # 添加这一行
            
            # 计算预测资产
            current_assets = df['total_assets'].iloc[-1]
            predicted_assets = current_assets * (1 + avg_growth_rate) ** prediction_period
            
            # 判断是否达到100万+
            target_reached = int(predicted_assets >= 1000000)
            
            # 特征工程 - 使用01编码处理分类变量
            features_for_clustering = [
                'age', 'monthly_income', 'total_assets', 'deposit_balance', 
                'financial_balance', 'fund_balance', 'insurance_balance',
                'product_count', 'financial_repurchase_count', 
                'credit_card_monthly_expense', 'investment_monthly_count',
                'app_login_count', 'app_financial_view_time', 
                'app_product_compare_count', 'avg_growth_rate'
            ]
            
            # 分类特征（只使用查询中包含的字段）
            categorical_features = ['gender', 'occupation']
            
            # 处理数值特征
            numeric_data = df[features_for_clustering].fillna(0)
            
            # 处理分类特征 - 使用01编码
            categorical_data = df[categorical_features].copy()
            label_encoders = {}
            
            from sklearn.preprocessing import LabelEncoder
            for feature in categorical_features:
                le = LabelEncoder()
                categorical_data[feature] = le.fit_transform(categorical_data[feature].astype(str))
                label_encoders[feature] = le
            
            # 合并所有特征
            X = pd.concat([numeric_data, categorical_data], axis=1)
            
            # 简化预测逻辑，不使用机器学习模型
            # 基于历史增长率进行简单预测
            if len(df) < 2:
                return f"""
                客户ID: {customer_id}
                客户姓名: {df['name'].iloc[-1] if 'name' in df.columns else '未知'}
                当前总资产: {current_assets:.2f}元
                
                由于历史数据不足（少于2个月），无法进行准确预测。
                建议：需要更多历史数据来进行资产增长预测。
                """
            
            # 基于增长率的简单预测
            recent_growth = df['growth_rate'].iloc[-3:].mean() if len(df) >= 3 else avg_growth_rate
            if np.isnan(recent_growth):
                recent_growth = 0
                
            # 保守预测和乐观预测
            conservative_prediction = current_assets * (1 + max(0, recent_growth * 0.5)) ** prediction_period
            optimistic_prediction = current_assets * (1 + max(0, recent_growth * 1.5)) ** prediction_period
            
            # 判断达到100万的概率
            if current_assets >= 1000000:
                probability = "已达到"
            elif optimistic_prediction >= 1000000:
                if conservative_prediction >= 1000000:
                    probability = "很高 (>80%)"
                else:
                    probability = "中等 (40-80%)"
            else:
                probability = "较低 (<40%)"
            
            # 可视化系数
            plt.figure(figsize=(10, 8))
            colors = ['blue' if x > 0 else 'red' for x in coef_df['coefficient']]
            plt.barh(range(len(coef_df)), coef_df['coefficient'], color=colors)
            plt.yticks(range(len(coef_df)), coef_df['feature'])
            plt.xlabel('Coefficient Value')
            plt.title('Logistic Regression Coefficients for Predicting 1M+ Assets in {} Months'.format(prediction_period))
            plt.axvline(x=0, color='black', linewidth=0.5)
            
            # 添加图例
            red_patch = plt.Rectangle((0,0),1,1, color='red')
            blue_patch = plt.Rectangle((0,0),1,1, color='blue')
            plt.legend([blue_patch, red_patch], ['Positive Coefficient', 'Negative Coefficient'])
            
            plt.tight_layout()
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'lr_coefficients_{customer_id}_{prediction_period}.png'
            save_path = os.path.join(save_dir, filename)
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()
            
            # 生成结果
            result = f"""
            客户ID: {customer_id}
            客户姓名: {df['name'].iloc[-1]}
            当前总资产: {current_assets:.2f}
            预测{prediction_period}个月后总资产: {predicted_assets:.2f}
            达到100万+概率: {target_reached}
            
            逻辑回归系数:
            {coef_df.to_markdown(index=False)}
            
            系数可视化图表已保存到 {filename}
            """
            
            return result
            
        except Exception as e:
            return f"预测执行失败: {str(e)}"

# ====== 初始化银行客户助手服务 ======
def init_agent_service():
    """初始化银行客户助手服务"""
    llm_cfg = {
        'model': 'qwen-plus-latest',
        'timeout': 30,
        'retry_count': 3,
    }
    tools = [
        {
            "mcpServers": {
                "tavily-mcp": {
                    # Tavily MCP配置
                    "command": "npx",
                    "args": ["-y", "tavily-mcp@0.1.4"],
                    "env": {
                        "TAVILY_API_KEY": os.getenv('TAVILY_API_KEY', ''),
                    },
                    "disabled": False,
                    "autoApprove": []
                },
            }
        },
        'exc_sql','lr_predict'
    ]
    try:
        bot = Assistant(
            llm=llm_cfg,
            name='百万客群经营助手',
            description='银行客户数据查询与分析',
            system_message=system_prompt,
            function_list=tools,  # 确保工具已注册
           # files=['./faq.txt']
        )
        print("助手初始化成功！")
        return bot
    except Exception as e:
        print(f"助手初始化失败: {str(e)}")
        raise


# ====== 定义 app_gui 函数以启动 Web 界面 ======
def app_gui():
    """图形界面模式，提供 Web 图形界面"""
    try:
        print("正在启动 Web 界面...")
        # 初始化助手
        bot = init_agent_service()
        # 配置聊天界面，列举典型查询问题
        chatbot_config = {
            'prompt.suggestions': [
                '我行目前有多少客户？总资产管理规模是多少？',
                '客户的平均资产是多少？高净值客户的占比如何？',
                '客户年龄分布情况如何？哪个年龄段客户最多？',
                '客户持有的四类产品（存款、理财、基金、保险）分布情况如何？',
                '根据历史数据，未来3个月的该银行的客户资产总规模统计会如何变化？',
                '请预测当前资产在80万到100万之间的客户中，有多少人未来3个月可能成为百万客户？',
                '316d72c0795f4fa9a61aeb8804c21b8f这个客户未来三个月是否有可能资产达到100万？'

            ]
        }
        print("Web 界面准备就绪，正在启动服务...")
        # 启动 Web 界面
        WebUI(
            bot,
            chatbot_config=chatbot_config
        ).run()
    except Exception as e:
        print(f"启动 Web 界面失败: {str(e)}")
        print("请检查网络连接和 API Key 配置")


if __name__ == '__main__':
    # 运行模式选择
    app_gui()          # 图形界面模式（默认）